import math
import struct
import time

import cv2
import numpy as np

from app.MspHelper import MspHelper
from myserial import CommService


class RobustOpticalFlow:
    def __init__(self, serial, camera_index=0):
        # 初始化摄像头
        self.cap = cv2.VideoCapture(camera_index)
        if not self.cap.isOpened():
            raise Exception("无法打开摄像头")

        # 设置摄像头参数（根据你的摄像头调整）
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 320)  # 降低分辨率以提高处理速度
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 240)
        self.cap.set(cv2.CAP_PROP_FPS, 60)


        # LK光流参数
        self.lk_params = dict(
            winSize=(15, 15),  # 搜索窗口大小
            maxLevel=2,  # 金字塔层数
            criteria=(cv2.TERM_CRITERIA_EPS | cv2.TERM_CRITERIA_COUNT, 10, 0.03)
        )

        # 特征点管理参数
        self.max_points = 150  # 最大跟踪点数
        self.min_points = 20  # 触发重新检测的最小点数
        self.detection_interval = 15  # 重新检测特征点的帧间隔
        self.frame_counter = 0

        # 状态变量
        self.old_gray = None
        self.p0 = None  # 上一帧的特征点
        self.initialized = False

        # 用于可视化（可选）
        self.visualize = True

        self.mspHelper = MspHelper()

        self.serial =serial

        print("光流模块初始化完成")

    def detect_features(self, frame, max_points=None):
        """在全图中检测特征点"""
        if max_points is None:
            max_points = self.max_points

        try:
            # 使用Shi-Tomasi角点检测
            features = cv2.goodFeaturesToTrack(
                frame,
                mask=None,
                maxCorners=max_points,
                qualityLevel=0.8,  # 降低质量要求以获取更多点
                minDistance=7,  # 特征点之间的最小距离
                blockSize=7
            )
        except Exception as ex:
            return None

        return features

    def calculate_flow_quality(self, status, flow_vectors):
        """计算光流数据的质量指标"""
        # 1. 基于成功跟踪点的比例 (0-255)
        if len(status) == 0:
            success_rate = 0
        else:
            success_rate = np.sum(status) / len(status)
        quality_from_success = int(success_rate * 255)

        # 2. 基于流向量的一致性
        if len(flow_vectors) > 1:
            flow_std = np.std(flow_vectors, axis=0)
            # 计算平均标准差并映射到质量值
            mean_std = np.mean(flow_std)
            consistency = np.exp(-0.1 * mean_std)  # 标准差越大，质量越低
            quality_from_consistency = int(consistency * 255)
        else:
            quality_from_consistency = 0

        # 综合质量（取两者中较小的，更保守）
        final_quality = min(quality_from_success, quality_from_consistency)

        return final_quality, success_rate


    def get_distance_measurement(self):
        """获取距离测量值（需要根据你的实际硬件实现）"""
        # 这里应该从你的激光测距仪或超声波传感器获取数据
        # 返回单位：米
        # 示例：返回固定值1米（仅用于测试）
        return 1.0

    def run(self):
        """主循环"""
        print("开始光流计算，按ESC退出")

        while True:
            # 读取帧
            ret, frame = self.cap.read()
            if not ret:
                print("无法读取帧，退出")
                break

            # 转换为灰度图
            frame_gray = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)

            # 初始化或重新初始化
            if not self.initialized or self.old_gray is None or self.p0 is None:
                self.p0 = self.detect_features(frame_gray)
                if self.p0 is not None and len(self.p0) > 5:
                    self.old_gray = frame_gray.copy()
                    self.initialized = True
                    print(f"初始化成功，检测到 {len(self.p0)} 个特征点")
                else:
                    print("初始化失败，无法检测到足够特征点")
                    time.sleep(0.1)
                    continue

            # 计算光流
            p1, st, err = cv2.calcOpticalFlowPyrLK(
                self.old_gray, frame_gray, self.p0, None, **self.lk_params
            )

            # 检查跟踪状态
            if p1 is None or st is None or np.sum(st) == 0:
                # 所有特征点丢失，紧急重新检测
                print("所有特征点丢失！紧急重新检测...")
                self.p0 = self.detect_features(frame_gray, max_points=self.max_points)
                if self.p0 is None or len(self.p0) == 0:
                    print("无法检测到任何特征点，发送无效数据")
                    # 发送无效数据（质量为0）
                    distance = self.get_distance_measurement()
                    # self.send_msp_optical_flow(0, 0, 0, distance)

                    # 重置状态，下一帧尝试重新初始化
                    self.initialized = False
                else:
                    # 更新旧帧，使用新检测的点
                    self.old_gray = frame_gray.copy()
                    print(f"紧急重新检测到 {len(self.p0)} 个特征点")

                # 跳过本次循环的后续处理
                continue

            # 提取成功跟踪的点
            good_new = p1[st == 1]
            good_old = self.p0[st == 1]

            # 计算光流向量
            flow_vectors = good_new - good_old

            # 计算统计数据
            mean_flow_x = np.mean(flow_vectors[:, 0])
            mean_flow_y = np.mean(flow_vectors[:, 1])

            # 计算质量指标
            quality, success_rate = self.calculate_flow_quality(st, flow_vectors)

            # 更新跟踪点列表
            self.p0 = good_new.reshape(-1, 1, 2)

            # 动态维护特征点
            self.frame_counter += 1
            current_points = len(self.p0)

            # 检查是否需要补充特征点
            if current_points < self.min_points or self.frame_counter >= self.detection_interval:
                print(f"特征点不足({current_points})或达到检测间隔，补充新点...")

                # 在全图检测新特征点
                new_features = self.detect_features(
                    frame_gray,
                    max_points=self.max_points - current_points
                )

                if new_features is not None and len(new_features) > 0:
                    # 合并新旧特征点
                    self.p0 = np.vstack((self.p0, new_features))
                    print(f"补充了 {len(new_features)} 个新特征点，当前总数: {len(self.p0)}")

                self.frame_counter = 0  # 重置计数器

            # 获取距离测量值
            distance = self.get_distance_measurement()

            # 通过MSP发送光流数据
            print(f"光流: ({mean_flow_x:.2f}, {mean_flow_y:.2f}), 质量: {quality}, 距离: {distance:.2f}m")


            # 发送MSP光流数据
            if self.serial is not None:
                multi10_x = round(mean_flow_x)
                multi10_y = round(mean_flow_y)
                flow_data = struct.pack('<Bii', quality, multi10_x, multi10_y)
                msp_frame = self.mspHelper.buildOneFrame(0x1F02, flow_data)
                self.serial.sendData(msp_frame)


            # 可视化（可选）
            if self.visualize:
                vis_frame = frame.copy()

                # 绘制跟踪点
                # for i, (new, old) in enumerate(zip(good_new, good_old)):
                #     a, b = new.ravel()
                #     c, d = old.ravel()
                #     # 绘制跟踪线
                #     vis_frame = cv2.line(vis_frame, (int(a), int(b)), (int(c), int(d)), (0, 255, 0), 2)
                #     # 绘制当前点
                #     vis_frame = cv2.circle(vis_frame, (int(a), int(b)), 3, (0, 0, 255), -1)

                # 显示状态信息
                cv2.putText(vis_frame, f"Points: {len(self.p0)}", (10, 20),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                cv2.putText(vis_frame, f"Quality: {quality}", (10, 40),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 255), 1)
                cv2.putText(vis_frame, f"Flow: ({mean_flow_x:.2f}, {mean_flow_y:.2f})", (10, 60),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, (255, 255, 0), 2)

                cv2.imshow('Optical Flow', vis_frame)

                # 按ESC退出
                if cv2.waitKey(1) & 0xFF == 27:
                    break

            # 更新前一帧
            self.old_gray = frame_gray.copy()

        # 清理资源
        self.cap.release()

        cv2.destroyAllWindows()
        print("光流模块已关闭")


if __name__ == '__main__':
    comm_service = CommService("/dev/ttyUSB0")

    obj = comm_service.open_comm()
    if obj is None:
        print('打开串口失败!')
    else:

        optical_flow = RobustOpticalFlow(comm_service,0)

        try:
            optical_flow.run()
        except KeyboardInterrupt:
            print("程序被用户中断")
        except Exception as e:
            print(f"程序运行出错: {e}")
